import torch
import torch.nn as nn
from .encoder import ResNet
from .bridge_network_cell import TF_model_cell
from .bridge_network_nerve import TF_model_nerve
# from networks.decoder_baseline import Decoder
from .decoder_cell_k import Decoder_K
# from networks.decoder_cell import Decoder
from .New_encoder import Encoder
from torchvision import models
# from networks.decoder_cell_k import Decoder_K
from .decoder_NFV3 import Decoder


class NF(nn.Module):
    def __init__(self):
        super(NF, self).__init__()
        self.encoder = Encoder()
        # self.model_cell = TF_model_cell()
        # self.model_nerve = TF_model_nerve()

        # self.decoder = Decoder_K()
        self.decoder = Decoder()

    def forward(self, input):
        feats = self.encoder(input)
        # for idx, feats in enumerate(feats):
        #     print("feats", feats[idx].shape)

        # feats_cell = self.model_cell(feats)
        # feats_nerve = self.model_nerve(feats)

        # out_cell = self.decoder_cell(feats_cell, feats)
        out = self.decoder(feats)
        # out = self.decoder(feats, feats)

        # out = self.decoder_cell(feats)

        return out
        # return out
